import json
import random
from pathlib import Path
from collections import Counter
import os

def load_jsonl(path):
    with open(path, 'r', encoding='utf-8') as f:
        return [json.loads(line.strip()) for line in f if line.strip()]

def save_jsonl(data, path):
    with open(path, 'w', encoding='utf-8') as f:
        for item in data:
            f.write(json.dumps(item, ensure_ascii=False) + '\n')

def main(input_path='data.jsonl', train_path='train.jsonl', test_path='test.jsonl', rate=0.9, seed=42):
    # 1. 读取数据
    data = load_jsonl(input_path)
    print(f"总数据量: {len(data)}")

    # 4. 随机打乱
    random.seed(seed)
    random.shuffle(data)

    # 5. 按 9:1 划分
    split_idx = int(rate * len(data))
    train_data = data[:split_idx]
    test_data = data[split_idx:]
    print(f"训练集数量: {len(train_data)}，测试集数量: {len(test_data)}")

    # 6. 保存结果
    save_jsonl(train_data, train_path)
    save_jsonl(test_data, test_path)
    print(f"已保存到 {train_path} 和 {test_path}")

if __name__ == '__main__':
    base_path = os.path.abspath(os.path.dirname(__file__))
    name = input()
    if name == "sard":
        input_path = os.path.join(base_path, 'data/sard_filtered_1000.jsonl')
        train_path = os.path.join(base_path, 'dataset/sard_train.jsonl')
        test_path = os.path.join(base_path, 'dataset/sard_test.jsonl')
    elif name == "sec1":
        input_path = os.path.join(base_path, 'data/sec1_all_filtered_2000.jsonl')
        train_path = os.path.join(base_path, 'dataset/sec1_train.jsonl')
        test_path = os.path.join(base_path, 'dataset/sec1_test.jsonl')
    elif name == "sec2":
        input_path = os.path.join(base_path, 'data/sec2_filtered_1000.jsonl')
        train_path = os.path.join(base_path, 'dataset/sec2_train.jsonl')
        test_path = os.path.join(base_path, 'dataset/sec2_test.jsonl')
    main(input_path=input_path, train_path=train_path, test_path=test_path)
